{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3f32b246-ad54-462e-9199-b837ebba3e67",
   "metadata": {},
   "source": [
    "# 天文221肖佳莹 实验报告"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "83ee3861-4e0a-418f-a939-3bdfe5838506",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import io\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import transforms, datasets\n",
    "import pyarrow.parquet as pq\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "from sklearn.model_selection import train_test_split\n",
    "from collections import Counter\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6bed9d04-90f4-40a6-a81d-58ab7e31bb96",
   "metadata": {},
   "outputs": [],
   "source": [
    "#数据\n",
    "# 读取Parquet文件\n",
    "table1 = pq.read_table(\"train-00000-of-00001.parquet\")\n",
    "df1 = table1.to_pandas()\n",
    "\n",
    "image_paths = df1['image']\n",
    "label1=df1[\"summary\"]\n",
    "# 处理标签\n",
    "class_names = [\"unknown\",\"smooth_round\",\"unbarred_spiral\",\"smooth_cigar\",\"edge_on_disk\",\"barred_spiral\",\n",
    "               \"featured_without_bar_or_spiral\"]\n",
    "class_to_idx = {name: idx for idx, name in enumerate(class_names)}\n",
    "valid_labels = []\n",
    "for label in df1['summary']:\n",
    "    if pd.isna(label) or (isinstance(label, str) and label.strip() == ''):\n",
    "        valid_labels.append(class_to_idx[\"unknown\"])\n",
    "    elif label in class_to_idx:\n",
    "        valid_labels.append(class_to_idx[label])\n",
    "    else:\n",
    "        valid_labels.append(class_to_idx[\"unknown\"])  # 未知类别归为 \"unknown\"\n",
    "\n",
    "labels = torch.tensor(valid_labels, dtype=torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fbc4146b-5955-4bf2-83d4-c4c1b1ec3a6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#训练集增强\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),  \n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a6a56106-9646-405a-84e5-00c390b19cce",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, image_paths, labels, transform=None):\n",
    "        self.image_paths = image_paths\n",
    "        self.labels = labels\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_paths)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # 从字节数据加载图像\n",
    "        img_bytes = self.image_paths.iloc[idx]['bytes'] \n",
    "        image = Image.open(io.BytesIO(img_bytes)).convert('RGB')\n",
    "        label = self.labels[idx]\n",
    "        \n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        return image, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "52c9aa2f-81c8-419f-ba21-79c6504a8e96",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 划分训练集和验证集\n",
    "train_data, val_data, train_labels, val_labels = train_test_split(\n",
    "    image_paths, labels, test_size=0.2, random_state=3575448\n",
    ")\n",
    "\n",
    "train_dataset = CustomDataset(train_data, train_labels,transform=transform)\n",
    "validation_dataset = CustomDataset(val_data, val_labels,transform=transform)\n",
    "\n",
    "# 创建DataLoader\n",
    "train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "validation_loader = DataLoader(validation_dataset, batch_size=64, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7063b0e4-4d71-40fd-bd88-1db2b28d0771",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "类别 | 出现次数\n",
      "-------------\n",
      "   0 |       2320\n",
      "   1 |        455\n",
      "   2 |        169\n",
      "   3 |        121\n",
      "   4 |         80\n",
      "   5 |         35\n",
      "   6 |         15\n"
     ]
    }
   ],
   "source": [
    "# 合并数据集并统计\n",
    "all_labels = torch.cat([train_labels, val_labels]).tolist()\n",
    "label_counts = Counter(all_labels)\n",
    "\n",
    "# 定义目标顺序的类别列表\n",
    "target_classes = list(range(7))  # [0, 1, 2, 3, 4, 5, 6]\n",
    "\n",
    "# 按顺序输出每个类别的出现次数\n",
    "print(\"类别 | 出现次数\")\n",
    "print(\"-------------\")\n",
    "for cls in target_classes:\n",
    "    count = label_counts.get(cls, 0)  # 如果类别不存在，返回0\n",
    "    print(f\"{cls:4d} | {count:10d}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "524c9e95-6394-4bb6-87fb-784cc0d5de41",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GalaxyCNN(nn.Module):\n",
    "    def __init__(self, num_classes=7):\n",
    "        super().__init__()\n",
    "        # 卷积层 + 批标准化\n",
    "        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(32)\n",
    "        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(64)\n",
    "        self.conv3 = ResidualBlock(64)   # 残差块\n",
    "        self.conv4 = ResidualBlock(64)   # 保持通道数一致\n",
    "        \n",
    "        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.global_avg_pool = nn.AdaptiveAvgPool2d((1, 1))  # 替代自适应池化\n",
    "        \n",
    "        # 全连接层 + 正则化\n",
    "        self.fc = nn.Linear(64, num_classes)  # 输入维度=64（与最后一层卷积输出一致）\n",
    "        self.dropout1 = nn.Dropout(0.3)\n",
    "        self.dropout2 = nn.Dropout(0.5)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.pool(F.relu(self.bn1(self.conv1(x))))\n",
    "        x = self.pool(F.relu(self.bn2(self.conv2(x))))\n",
    "        x = self.conv3(x)  # 残差块\n",
    "        x = self.conv4(x)  # 残差块\n",
    "        x = self.global_avg_pool(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.dropout1(x)\n",
    "        x = self.fc(x)\n",
    "        return x\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, in_channels):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(in_channels)\n",
    "        self.conv2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(in_channels)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "        x = F.relu(self.bn1(self.conv1(x)))\n",
    "        x = self.bn2(self.conv2(x))\n",
    "        x += residual\n",
    "        return F.relu(x)\n",
    "\n",
    "model = GalaxyCNN(num_classes=7).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d5413820-f8ab-4e53-8c4e-871c3cbde6bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class_weights = torch.tensor([3.5, 8.0, 12.5, 15.0, 18.0, 20.0, 23.0]).to(device)\n",
    "criterion = nn.CrossEntropyLoss(weight=class_weights)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, weight_decay=1e-4)\n",
    "scheduler = StepLR(optimizer, step_size=5, gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "21cdd60d-7714-41d4-acb6-db6407795c24",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.15it/s]\n",
      "Epoch 1 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1:\n",
      "Train Loss: 1.8536, Train Acc: 34.66%\n",
      "Val Loss: 1.8489, Val Acc: 14.87%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.22it/s]\n",
      "Epoch 2 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.51it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2:\n",
      "Train Loss: 1.5960, Train Acc: 58.72%\n",
      "Val Loss: 1.4647, Val Acc: 54.93%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 3 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.21it/s]\n",
      "Epoch 3 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3:\n",
      "Train Loss: 1.5279, Train Acc: 63.97%\n",
      "Val Loss: 1.4164, Val Acc: 59.47%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 4 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.22it/s]\n",
      "Epoch 4 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4:\n",
      "Train Loss: 1.5145, Train Acc: 61.58%\n",
      "Val Loss: 1.4223, Val Acc: 61.50%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 5 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.22it/s]\n",
      "Epoch 5 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5:\n",
      "Train Loss: 1.4944, Train Acc: 63.22%\n",
      "Val Loss: 1.3668, Val Acc: 67.76%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 6 [Train]: 100%|██████████| 40/40 [00:17<00:00,  2.23it/s]\n",
      "Epoch 6 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.53it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 6:\n",
      "Train Loss: 1.4618, Train Acc: 62.05%\n",
      "Val Loss: 1.3617, Val Acc: 68.08%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 7 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.22it/s]\n",
      "Epoch 7 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 7:\n",
      "Train Loss: 1.4619, Train Acc: 62.75%\n",
      "Val Loss: 1.3706, Val Acc: 69.80%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 8 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.22it/s]\n",
      "Epoch 8 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.51it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 8:\n",
      "Train Loss: 1.4570, Train Acc: 64.24%\n",
      "Val Loss: 1.3666, Val Acc: 69.33%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 9 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.22it/s]\n",
      "Epoch 9 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.51it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9:\n",
      "Train Loss: 1.4584, Train Acc: 63.62%\n",
      "Val Loss: 1.3660, Val Acc: 69.48%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 10 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.21it/s]\n",
      "Epoch 10 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.55it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10:\n",
      "Train Loss: 1.4409, Train Acc: 64.87%\n",
      "Val Loss: 1.3681, Val Acc: 69.95%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 11 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.21it/s]\n",
      "Epoch 11 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 11:\n",
      "Train Loss: 1.4464, Train Acc: 64.05%\n",
      "Val Loss: 1.3684, Val Acc: 69.95%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 12 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.21it/s]\n",
      "Epoch 12 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.49it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 12:\n",
      "Train Loss: 1.4477, Train Acc: 64.36%\n",
      "Val Loss: 1.3685, Val Acc: 69.95%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 13 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.20it/s]\n",
      "Epoch 13 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 13:\n",
      "Train Loss: 1.4687, Train Acc: 63.89%\n",
      "Val Loss: 1.3683, Val Acc: 69.95%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 14 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.22it/s]\n",
      "Epoch 14 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 14:\n",
      "Train Loss: 1.4597, Train Acc: 64.12%\n",
      "Val Loss: 1.3675, Val Acc: 69.95%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 15 [Train]: 100%|██████████| 40/40 [00:18<00:00,  2.22it/s]\n",
      "Epoch 15 [Val]: 100%|██████████| 10/10 [00:02<00:00,  3.48it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 15:\n",
      "Train Loss: 1.4381, Train Acc: 65.45%\n",
      "Val Loss: 1.3666, Val Acc: 69.64%\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 15\n",
    "train_losses = []\n",
    "train_accuracies = []\n",
    "val_losses = []\n",
    "val_accuracies = []\n",
    "train_predictions = []\n",
    "val_predictions = []\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    train_correct = 0\n",
    "    train_total = 0\n",
    "    train_running_loss = 0.0\n",
    "    epoch_train_preds = []\n",
    "    \n",
    "    # 训练阶段（同时计算训练集准确率）\n",
    "    for images, labels in tqdm(train_loader, desc=f\"Epoch {epoch+1} [Train]\"):\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        # 计算训练集准确率\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        epoch_train_preds.extend(predicted.cpu().numpy())\n",
    "        \n",
    "        train_total += labels.size(0)\n",
    "        train_correct += (predicted == labels).sum().item()\n",
    "        train_running_loss += loss.item()\n",
    "        \n",
    "    train_predictions.append(epoch_train_preds)\n",
    "    \n",
    "    # 计算训练集指标\n",
    "    train_loss = train_running_loss / len(train_loader)\n",
    "    train_acc = 100 * train_correct / train_total\n",
    "    train_losses.append(train_loss)\n",
    "    train_accuracies.append(train_acc)\n",
    "    \n",
    "    # 验证阶段\n",
    "    model.eval()\n",
    "    val_correct = 0\n",
    "    val_total = 0\n",
    "    val_running_loss = 0.0\n",
    "    epoch_val_preds = []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for images, labels in tqdm(validation_loader, desc=f\"Epoch {epoch+1} [Val]\"):\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels) \n",
    "            val_running_loss += loss.item()\n",
    "            \n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            epoch_val_preds.extend(predicted.cpu().numpy())\n",
    "            \n",
    "            val_total += labels.size(0)\n",
    "            val_correct += (predicted == labels).sum().item()\n",
    "\n",
    "    val_predictions.append(epoch_val_preds)\n",
    "    \n",
    "    # 计算验证集指标\n",
    "    val_loss = val_running_loss / len(validation_loader)\n",
    "    val_acc = 100 * val_correct / val_total\n",
    "    val_losses.append(val_loss)\n",
    "    val_accuracies.append(val_acc)\n",
    "      \n",
    "    # 输出结果\n",
    "    print(f\"Epoch {epoch+1}:\\n\"\n",
    "          f\"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%\\n\"\n",
    "          f\"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%\")\n",
    "    print(\"-\" * 50)\n",
    "\n",
    "    # 更新学习率\n",
    "    scheduler.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1af521e2-a6d4-4046-b785-c181605adfea",
   "metadata": {},
   "outputs": [],
   "source": [
    "converted_last1 = [int(x) for x in train_predictions[-1]]\n",
    "#print(\"最后一个epoch转换结果:\", converted_last1)\n",
    "converted_last2 = [int(x) for x in val_predictions[-1]]\n",
    "#print(\"最后一个epoch转换结果:\", converted_last2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1702a0da-4c7c-44d4-bcdd-d9ca5a286acb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "模型预测\n",
      "类别 | 出现次数\n",
      "-------------\n",
      "   0 |       2664\n",
      "   1 |        329\n",
      "   2 |        167\n",
      "   3 |         19\n",
      "   4 |         14\n",
      "   5 |          2\n",
      "   6 |          0\n"
     ]
    }
   ],
   "source": [
    "total_all=converted_last1+converted_last2\n",
    "result = Counter(total_all)\n",
    "# 定义目标顺序的类别列表（0到6）\n",
    "target_classes1 = list(range(7))  # [0, 1, 2, 3, 4, 5, 6]\n",
    "\n",
    "# 按顺序输出每个类别的出现次数\n",
    "print(\"模型预测\")\n",
    "print(\"类别 | 出现次数\")\n",
    "print(\"-------------\")\n",
    "for cls in target_classes1:\n",
    "    count = result.get(cls, 0)  # 如果类别不存在重复，显示0\n",
    "    print(f\"{cls:4d} | {count:10d}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b8d9f7f7-5e16-483b-9f19-fafa5182678f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存模型参数\n",
    "torch.save(model.state_dict(), \"galaxy.pth\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fc76d25-26da-4c4f-ad63-e7e53918e57e",
   "metadata": {},
   "source": [
    "# 在测试集上运行模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d1bac8ff-2a41-47aa-849e-f2291a8ac84c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GalaxyCNN(\n",
       "  (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv3): ResidualBlock(\n",
       "    (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (conv4): ResidualBlock(\n",
       "    (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  (global_avg_pool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (fc): Linear(in_features=64, out_features=7, bias=True)\n",
       "  (dropout1): Dropout(p=0.3, inplace=False)\n",
       "  (dropout2): Dropout(p=0.5, inplace=False)\n",
       ")"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#测试集\n",
    "# 1. 加载模型\n",
    "model = GalaxyCNN(num_classes=7)\n",
    "model.load_state_dict(torch.load('galaxy.pth'))\n",
    "model = model.to(device)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3484899a-d5a4-4447-8e3e-fda4f413f796",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. 准备测试集\n",
    "#数据\n",
    "# 读取Parquet文件\n",
    "table2 = pq.read_table(\"test-00000-of-00001.parquet\")\n",
    "df2 = table2.to_pandas()\n",
    "\n",
    "image_paths2 = df2['image']\n",
    "label1=df2[\"summary\"]\n",
    "# 处理标签\n",
    "class_names = [\"unknown\",\"smooth_round\",\"unbarred_spiral\",\"smooth_cigar\",\"edge_on_disk\",\"barred_spiral\",\n",
    "               \"featured_without_bar_or_spiral\"]\n",
    "class_to_idx = {name: idx for idx, name in enumerate(class_names)}\n",
    "valid_labels1 = []\n",
    "for label in label1 :\n",
    "    if pd.isna(label) or (isinstance(label, str) and label.strip() == ''):\n",
    "        valid_labels1.append(class_to_idx[\"unknown\"])\n",
    "    elif label in class_to_idx:\n",
    "        valid_labels1.append(class_to_idx[label])\n",
    "    else:\n",
    "        valid_labels1.append(class_to_idx[\"unknown\"])  # 未知类别归为 \"unknown\"\n",
    "\n",
    "labels_t = torch.tensor(valid_labels1, dtype=torch.long)\n",
    "\n",
    "test_dataset = CustomDataset(image_paths2, labels_t, transform=transform)\n",
    "test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a216862-3343-463b-986a-9381589ed06c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5999ec40-9b50-4201-b504-bbf45578a4d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.7293\n"
     ]
    }
   ],
   "source": [
    "# 3. 预测\n",
    "# 初始化变量\n",
    "test_correct = 0\n",
    "test_total = 0\n",
    "\n",
    "with torch.no_grad():\n",
    "    for images, labels in test_loader:\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        outputs = model(images)\n",
    "        _, preds = torch.max(outputs, 1)\n",
    "        \n",
    "        # 更新正确预测和总预测的数量\n",
    "        test_correct += (preds == labels).sum().item()\n",
    "        test_total += labels.size(0)\n",
    "\n",
    "# 计算准确率\n",
    "accuracy = test_correct / test_total if test_total != 0 else 0\n",
    "\n",
    "print(f\"Accuracy: {accuracy:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "814cf9b9-28ee-47f0-aad0-cd0e4a9370c2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff6b8472-6c14-496d-80c3-db9c0b0931ef",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33f164e2-91f3-4f0d-a63f-85dcd51fc30e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
